global numBins

classLabels = ['A' 'B' 'C' 'D' 'F' 'G' 'H' 'K' 'M' 'O' 'P' 'S' 'Z']; % note: currently do not use W - water bottle
C = numel(classLabels);
trainPath = './images/training/';
masksPath = './images/trainingmasks/';

testImgsPath = './images/test/easy/';
colorSpace = 'rgb';
distanceComparison = 'cosine';
similarity = 'weightedComponents';

numBins = 100;

% Build the training models
if ~exist(['./data/models_' colorSpace '_' num2str(numBins) '.mat'])
    models = -ones(5,numBins,C);
    for i=1:C
        classPath = [trainPath classLabels(i) '/'];
        list = dir([classPath '*.jpg']);
        list = {list(:).name};
        % For all training images of class i, build model color and texture
        % histograms
        Chist1 = zeros(1,numBins); Chist2 = zeros(1,numBins); Chist3 = zeros(1,numBins);
        Thist1 = zeros(1,numBins); Thist2 = zeros(1,numBins);
        for j=1:numel(list)
            imgName = list{j};
            I = imread([classPath imgName]);
            mask = load([masksPath imgName(1:end-4)]);
            
            [chist1,chist2,chist3] = Color_Features(I,mask.mask,colorSpace,false);
            [thist1,thist2] = Texture_Features(mask.mask,I,0);
            Chist1 = Chist1 + chist1; Chist2 = Chist2 + chist2; Chist3 = Chist3 + chist3;
            Thist1 = Thist1 + thist1; Thist2 = Thist2 + thist2;
        end
        % Normalize the class histograms
        models(1,:,i) = Chist1/sum(Chist1);
        models(2,:,i) = Chist2/sum(Chist2);
        models(3,:,i) = Chist3/sum(Chist3);
        models(4,:,i) = Thist1/sum(Thist1);
        models(5,:,i) = Thist2/sum(Thist2);
        disp(['Models built for class ' classLabels(i) '.'])
    end
    save(['./data/models_' colorSpace '_' num2str(numBins)], 'models')
else
    load(['./data/models_' colorSpace '_' num2str(numBins)])
end

% Build the training mean HS models to be used for segmentation
if ~exist('./data/models_meansHS.mat')
    models_meansHS = -ones(C,2);
    for i=1:C
        classPath = [trainPath classLabels(i) '/'];
        list = dir([classPath '*.jpg']);
        list = {list(:).name};
        
        sumPixels = 0;
        sumHue = 0;
        sumSat = 0;
        for j=1:numel(list)
            imgName = list{j};
            I = imread([classPath imgName]);
            mask = load([masksPath imgName(1:end-4)]); mask = mask.mask;
            hsv = rgb2hsv(I);
            H = hsv(:,:,1);
            S = hsv(:,:,2);
            sumPixels = sumPixels + sum(mask(:));
            sumHue = sumHue + sum(H(mask));
            sumSat = sumSat + sum(S(mask));
        end
        models_meansHS(i,1) = sumHue/sumPixels;
        models_meansHS(i,2) = sumSat/sumPixels;
        disp(['Mean hue and saturation saved for class ' classLabels(i) '.'])
    end
    save('./data/models_meansHS.mat', 'models_meansHS');
else
    load('./data/models_meansHS.mat')
end

% Make a mask of true classes for the test images
list = dir([testImgsPath '*.jpg']); list = {list(:).name};
T = numel(list); % T = number of test images
testClasses = zeros(1,T);
for i=1:numel(list)
    parts = textscan(list{i}, '%s%s%s', 'Delimiter', '-.');
    testClasses(i) = strfind(classLabels, char(parts{1}));
end
truthMask = zeros(C,T);
truthMask(sub2ind(size(truthMask), testClasses, 1:T)) = 1;

testEvalMat = load('./testEvalMat_easy.mat');

for t=1:T
    testFile = [testImgsPath list{t}];
    I = imread(testFile);
    % Get mask(s) from segmentation results
    ROIs = segmentationL(I, models_meansHS);
    
%     subplot(1,size(ROIs,3)+1,1)
%     imshow(I)
    
    testX = zeros(T,numBins*5);
    testF = zeros(1,numBins*5);
    testFClasses = 0;
    for r=1:size(ROIs,3)
        mask = logical(ROIs(:,:,r));
        
%         subplot(1,size(ROIs,3)+1,r+1)
%         imshow(mask)
        
        [chist1,chist2,chist3] = Color_Features(I,mask,colorSpace,true);
        [thist1,thist2] = Texture_Features(mask,I,1);
        x = horzcat(chist1,chist2,chist3,thist1,thist2);
        
        trueLoc = testEvalMat.locations(strmatch(list{t},testEvalMat.filenames,'exact'),:);
        trueLoc = round(trueLoc);
        if mask(sub2ind(size(mask),trueLoc(1),trueLoc(2)))==1
            testX(t,:) = x;
        else 
            testF = vertcat(testF,x);
            testFClasses = horzcat(testClasses(1,t));
        end
    end
    disp(['Extracted features for test image ' num2str(t) ' of ' num2str(T) '.'])
end
testF = testF(2:end,:);
testFClasses = testFClasses(1,2:end);


% Classify each test image
scoresX = zeros(C,T);
tic;
for t=1:T
 
    D = zeros(size(models,3),5);
    sumD = zeros(size(models,3),1);
    
    % ----------------------------------------------------------------
    % Calculate similarity of test histogram(s) with the model
    % histogram(s) from each class and assign to class with highest
    % similarity
    chist1 = testX(t,1:numBins);
    chist2 = testX(t,numBins+1:2*numBins);
    chist3 = testX(t,2*numBins+1:3*numBins);
    thist1 = testX(t,3*numBins+1:4*numBins);
    thist2 = testX(t,4*numBins+1:5*numBins);
    
    vector = horzcat(chist1,chist2,chist3,thist1,thist2);
    P = models;
    for f = 1:size(models,3)
        if strcmp(distanceComparison,'manhattan')
            D(f,1) = pdist2(chist1,P(1,:,f),'cityblock');
            D(f,2) = pdist2(chist2,P(2,:,f),'cityblock');
            D(f,3) = pdist2(chist3,P(3,:,f),'cityblock');
            D(f,4) = pdist2(thist1,P(4,:,f),'cityblock');
            D(f,5) = pdist2(thist2,P(5,:,f),'cityblock');
        elseif strcmp(distanceComparison,'cosine')
            D(f,1) = pdist2(chist1,P(1,:,f),'cosine');
            D(f,2) = pdist2(chist2,P(2,:,f),'cosine');
            D(f,3) = pdist2(chist3,P(3,:,f),'cosine');
            D(f,4) = pdist2(thist1,P(4,:,f),'cosine');
            D(f,5) = pdist2(thist2,P(5,:,f),'cosine');
        elseif strcmp(distanceComparison,'intersection')
            D(f,1) = 1-sum(min(chist1,P(1,:,f)));
            D(f,2) = 1-sum(min(chist2,P(2,:,f)));
            D(f,3) = 1-sum(min(chist3,P(3,:,f)));
            D(f,4) = 1-sum(min(thist1,P(4,:,f)));
            D(f,5) = 1-sum(min(thist2,P(5,:,f)));
        end
    end
    
    sumD = sum(D(:),2);
    scoresX(:,t) = sumD;
end

F = size(testF,1);
scoresF = zeros(C,F);
for t=1:F
 
    D = zeros(size(models,3),5);
    sumD = zeros(size(models,3),1);
    
    % ----------------------------------------------------------------
    % Calculate similarity of test histogram(s) with the model
    % histogram(s) from each class and assign to class with highest
    % similarity
    chist1 = testF(t,1:numBins);
    chist2 = testF(t,numBins+1:2*numBins);
    chist3 = testF(t,2*numBins+1:3*numBins);
    thist1 = testF(t,3*numBins+1:4*numBins);
    thist2 = testF(t,4*numBins+1:5*numBins);
    
    vector = horzcat(chist1,chist2,chist3,thist1,thist2);
    P = models;
    for f = 1:size(models,3)
        if strcmp(distanceComparison,'manhattan')
            D(f,1) = pdist2(chist1,P(1,:,f),'cityblock');
            D(f,2) = pdist2(chist2,P(2,:,f),'cityblock');
            D(f,3) = pdist2(chist3,P(3,:,f),'cityblock');
            D(f,4) = pdist2(thist1,P(4,:,f),'cityblock');
            D(f,5) = pdist2(thist2,P(5,:,f),'cityblock');
        elseif strcmp(distanceComparison,'cosine')
            D(f,1) = pdist2(chist1,P(1,:,f),'cosine');
            D(f,2) = pdist2(chist2,P(2,:,f),'cosine');
            D(f,3) = pdist2(chist3,P(3,:,f),'cosine');
            D(f,4) = pdist2(thist1,P(4,:,f),'cosine');
            D(f,5) = pdist2(thist2,P(5,:,f),'cosine');
        elseif strcmp(distanceComparison,'intersection')
            D(f,1) = 1-sum(min(chist1,P(1,:,f)));
            D(f,2) = 1-sum(min(chist2,P(2,:,f)));
            D(f,3) = 1-sum(min(chist3,P(3,:,f)));
            D(f,4) = 1-sum(min(thist1,P(4,:,f)));
            D(f,5) = 1-sum(min(thist2,P(5,:,f)));
        end
    end
    
    sumD = sum(D(:),2);
    scoresF(:,t) = sumD;
end
time = toc;

% Calculate accuracy
[maxScoresX, classIds] = sort(scoresX,'ascend');
rank = 1;
accuracy = sum(truthMask(sub2ind(size(truthMask), ids(rank,:), 1:T)) > 0) / T;
correct = truthMask(sub2ind(size(truthMask), ids(rank,:), 1:T)) > 0;
incorrect = truthMask(sub2ind(size(truthMask), ids(rank,:), 1:T)) < 0;

[maxScoresF, classIdsF] = sort(scoresF,'ascend');

disp(['Accuracy for [' distanceComparison ', ' colorSpace '] is ' num2str(accuracy*100) '%.'])
disp(['Segmentation detected ' num2str(F) ' regions that are not food items.'])
disp(['The average score of a miscellaneous regions detected is ' mean(maxScoresF(1,:)) '.'])
disp(['Average time to classify a test image is ' num2str(time/T) ' seconds.'])



